import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import copy

from common.Common import Point
from env.env import EdgeServer, Service, User, MECEnv

# hyper-parameters
BATCH_SIZE = 128
LR = 0.0005
GAMMA = 0.90
EPISILO = 0.9
MEMORY_CAPACITY = 1000
Q_NETWORK_ITERATION = 100 # copy targetnet的间隔


def make_env():
    # 一个user 三个server
    edge_server_list = []

    edge_server_list.append(EdgeServer(0, Point(500, 9), 30, 10))
    edge_server_list.append(EdgeServer(1, Point(2500, 9), 30, 10))
    edge_server_list.append(EdgeServer(2, Point(4500, 9), 30, 10))

    service = Service(50, 20)

    user_list = [[]]
    time_slot = 5000 / 500
    for i in range(500):
        user_list[0].append(User(0, Point((time_slot * i) % 5000, 9), service))

    return MECEnv(edge_server_list, user_list, 100)

env = make_env()
NUM_ACTIONS = 3
NUM_STATES = 6
ENV_A_SHAPE = 0


class Net(nn.Module):
    """docstring for Net"""

    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(NUM_STATES, 50)
        self.fc1.weight.data.normal_(0, 0.1)
        self.fc2 = nn.Linear(50, 30)
        self.fc2.weight.data.normal_(0, 0.1)
        self.out = nn.Linear(30, NUM_ACTIONS)
        self.out.weight.data.normal_(0, 0.1)

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = F.relu(x)
        action_prob = self.out(x)
        return action_prob


class DQN():
    """docstring for DQN"""

    def __init__(self):
        super(DQN, self).__init__()
        self.eval_net, self.target_net = Net(), Net()

        self.learn_step_counter = 0
        self.memory_counter = 0
        self.memory = np.zeros((MEMORY_CAPACITY, NUM_STATES * 2 + 2))  # num_states state的维度
        # When we store the memory, we put the state, action, reward and next_state in the memory
        # here reward and action is a number, state is a ndarray
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=LR)
        self.loss_func = nn.MSELoss()

    def choose_action(self, state):
        state = torch.unsqueeze(torch.FloatTensor(state), 0)  # get a 1D array
        if np.random.randn() <= EPISILO:  # greedy policy
            # 选择q值最大的action
            action_value = self.eval_net.forward(state)
            action = torch.max(action_value, 2)[1].data.numpy()
            action = action[0][0] if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE)
        else:  # random policy
            action = np.random.randint(3)
            # print(action)
            action = action if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE)
        return action

    def store_transition(self, state, action, reward, next_state):
        transition = np.hstack((state[0], [action, reward], next_state[0]))
        index = self.memory_counter % MEMORY_CAPACITY
        self.memory[index, :] = transition  # 新的经验
        self.memory_counter += 1

    def learn(self):

        # update the parameters
        if self.learn_step_counter % Q_NETWORK_ITERATION == 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
        self.learn_step_counter += 1

        # sample batch from memory
        sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE)  # 先随机出要采样的下标
        batch_memory = self.memory[sample_index, :] # 再从memory里取出采样的path
        batch_state = torch.FloatTensor(batch_memory[:, :NUM_STATES])
        batch_action = torch.LongTensor(batch_memory[:, NUM_STATES:NUM_STATES + 1].astype(int))
        batch_reward = torch.FloatTensor(batch_memory[:, NUM_STATES + 1:NUM_STATES + 2])
        batch_next_state = torch.FloatTensor(batch_memory[:, -NUM_STATES:])

        # q_eval
        # 关键算法步骤
        # q_eval_before = self.eval_net(batch_state)
        # q_eval = q_eval_before.gather(1, batch_action) # 查询得到动作对应的q_val 因为q_eval_before返回的是动作维度的返回值，所以选择一个动作对应的
        # q_next_before = self.target_net(batch_next_state)
        # q_next = q_next_before.detach() # 返回一个新的tensor，从当前计算图中分离下来的，但是仍指向原变量的存放位置,不同之处只是requires_grad为false，得到的这个tensor永远不需要计算其梯度，不具有grad。
        # q_target = batch_reward + GAMMA * q_next.max(1)[0].view(BATCH_SIZE, 1)
        # loss = self.loss_func(q_eval, q_target)

        q_eval = self.eval_net(batch_state).gather(1, batch_action)
        q_next = self.target_net(batch_next_state).detach()
        q_target = batch_reward + GAMMA * q_next.max(1)[0].view(BATCH_SIZE, 1)
        loss = self.loss_func(q_eval, q_target)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()



def main():
    dqn = DQN()
    episodes = 400
    print("Collecting Experience....")
    reward_list = []
    # plt.ion()
    fig, ax = plt.subplots()
    for i in range(episodes):
        state = env.reset()
        ep_reward = 0
        while True:
            action = dqn.choose_action(state)
            next_state, reward, done, _ = env.step(action)
            dqn.store_transition(state, action, reward, next_state)
            ep_reward += reward

            if dqn.memory_counter >= MEMORY_CAPACITY:
                dqn.learn()
                if done:
                    print("episode: {} , the episode reward is {}".format(i, ep_reward))
            if done:
                break
            state = next_state
        r = copy.copy(ep_reward)
        reward_list.append(r)
        # ax.set_xlim(0, 5000)
        # ax.cla()
        # ax.plot(reward_list, 'g-', label='total_loss')
        # plt.pause(0.001)
        # plt.show()
    ax.plot(reward_list)
    plt.show()


if __name__ == '__main__':
    main()
